use std::ops::{Add, Div, Mul, Neg, Sub};

#[derive(Clone, Copy)]
pub struct Variable {
    value: f64,
    grad: f64,
}

impl Variable {
    pub fn new(value: f64) -> Self {
        Self { value, grad: 0.0 }
    }
    pub fn new_diff(value: f64) -> Self {
        Self { value, grad: 1.0 }
    }
    pub fn grad(self) -> f64 {
        self.grad
    }
    pub fn value(self) -> f64 {
        self.value
    }
    pub fn value_grad(self) -> (f64, f64) {
        (self.value, self.grad)
    }
}

impl Add for Variable {
    type Output = Self;
    fn add(self, rhs: Self) -> Self::Output {
        Self {
            value: self.value + rhs.value,
            grad: self.grad + rhs.grad,
        }
    }
}

impl Add<f64> for Variable {
    type Output = Self;
    fn add(self, rhs: f64) -> Self::Output {
        Self {
            value: self.value + rhs,
            ..self
        }
    }
}

impl Sub for Variable {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self::Output {
        Self {
            value: self.value - rhs.value,
            grad: self.grad - rhs.grad,
        }
    }
}

impl Sub<f64> for Variable {
    type Output = Self;
    fn sub(self, rhs: f64) -> Self::Output {
        Self {
            value: self.value - rhs,
            ..self
        }
    }
}

impl Mul for Variable {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self::Output {
        Self {
            value: self.value * rhs.value,
            grad: self.value * rhs.grad + self.grad * rhs.value,
        }
    }
}

impl Mul<f64> for Variable {
    type Output = Self;
    fn mul(self, rhs: f64) -> Self::Output {
        Self {
            value: self.value * rhs,
            grad: self.grad * rhs,
        }
    }
}

impl Div for Variable {
    type Output = Self;
    fn div(self, rhs: Self) -> Self::Output {
        let value = self.value / rhs.value;
        Self {
            value,
            grad: value + (self.grad - value * rhs.grad) / rhs.value,
        }
    }
}

impl Div<f64> for Variable {
    type Output = Self;
    fn div(self, rhs: f64) -> Self::Output {
        Self {
            value: self.value / rhs,
            grad: self.grad / rhs,
        }
    }
}

impl Neg for Variable {
    type Output = Self;
    fn neg(self) -> Self::Output {
        Self {
            value: -self.value,
            grad: -self.grad,
        }
    }
}

impl Variable {
    pub fn recip(self) -> Self {
        let value = self.value.recip();
        Self {
            value,
            grad: -self.grad * value * value,
        }
    }

    pub fn sin(self) -> Self {
        Self {
            value: self.value.sin(),
            grad: self.value.cos() * self.grad,
        }
    }
    pub fn sinh(self) -> Self {
        Self {
            value: self.value.sinh(),
            grad: self.value.cosh() * self.grad,
        }
    }
    pub fn asin(self) -> Self {
        Self {
            value: self.value.asin(),
            grad: 1.0 / (1.0 - self.value.powi(2)).sqrt() * self.grad,
        }
    }
    pub fn asinh(self) -> Self {
        Self {
            value: self.value.asinh(),
            grad: 1.0 / (self.value.powi(2) + 1.0).sqrt() * self.grad,
        }
    }
    pub fn cos(self) -> Self {
        Self {
            value: self.value.cos(),
            grad: -self.value.sin() * self.grad,
        }
    }
    pub fn cosh(self) -> Self {
        Self {
            value: self.value.cosh(),
            grad: self.value.sinh() * self.grad,
        }
    }
    pub fn acos(self) -> Self {
        Self {
            value: self.value.acos(),
            grad: -1.0 / (1.0 - self.value.powi(2)).sqrt() * self.grad,
        }
    }
    pub fn acosh(self) -> Self {
        Self {
            value: self.value.asinh(),
            grad: 1.0 / (self.value.powi(2) - 1.0).sqrt() * self.grad,
        }
    }
    pub fn tan(self) -> Self {
        Self {
            value: self.value.tan(),
            grad: self.value.cos().powi(2).recip() * self.grad,
        }
    }
    pub fn tanh(self) -> Self {
        Self {
            value: self.value.tanh(),
            grad: self.grad / self.value.cosh().powi(2),
        }
    }
    pub fn atan(self) -> Self {
        Self {
            value: self.value.atan(),
            grad: self.grad / (1.0 + self.value.powi(2)),
        }
    }
    pub fn atanh(self) -> Self {
        Self {
            value: self.value.atanh(),
            grad: self.grad / (1.0 - self.value.powi(2)),
        }
    }

    pub fn exp(self) -> Self {
        let value = self.value.exp();
        Self {
            value,
            grad: value * self.grad,
        }
    }
    pub fn exp2(self) -> Self {
        let value = self.value.exp2();
        Self {
            value: value,
            grad: value * 2f64.ln() * self.grad,
        }
    }
    pub fn exp_m1(self) -> Self {
        Self {
            value: self.value.exp_m1(),
            grad: self.value.exp() * self.grad,
        }
    }

    pub fn ln(self) -> Self {
        Self {
            value: self.value.ln(),
            grad: self.grad / self.value,
        }
    }
    pub fn log10(self) -> Self {
        Self {
            value: self.value.log10(),
            grad: self.grad / (self.value * 10f64.ln()),
        }
    }
    pub fn log2(self) -> Self {
        Self {
            value: self.value.log2(),
            grad: self.grad / (self.value * 2f64.ln()),
        }
    }
    pub fn log(self, base: f64) -> Self {
        Self {
            value: self.value.log(base),
            grad: self.grad / (self.value * base.ln()),
        }
    }

    pub fn logx(self, base: Self) -> Self {
        let ln_b_recip = base.value.ln().recip();
        let value = self.value.log(base.value);
        Self {
            value,
            grad: ln_b_recip * self.value * self.grad + value * ln_b_recip * base.grad / base.value,
        }
    }

    pub fn powf(self, n: f64) -> Self {
        Self {
            value: self.value.powf(n),
            grad: n * self.value.powf(n - 1.0) * self.grad,
        }
    }
    pub fn powi(self, n: i32) -> Self {
        Self {
            value: self.value.powi(n),
            grad: n as f64 * self.value.powi(n - 1) * self.grad,
        }
    }
    pub fn powx(self, n: Self) -> Self {
        let value = self.value.powf(n.value);
        Self {
            value,
            grad: n.value * self.value.powf(n.value - 1.0) * self.grad
                + value * self.value.ln() * n.grad,
        }
    }
    pub fn sqrt(self) -> Self {
        let value = self.value.sqrt();
        Self {
            value,
            grad: self.grad / (2.0 * value),
        }
    }
    pub fn cbrt(self) -> Self {
        let value = self.value.cbrt();
        Self {
            value,
            grad: self.grad / (value.powi(2) * 3.0),
        }
    }
    pub fn hypot(self, other: Self) -> Self {
        let value = (self.value.powi(2) + other.value.powi(2)).sqrt();
        Self {
            value,
            grad: self.value / value * self.grad + other.value / value * other.grad,
        }
    }
}
